#!/usr/bin/env python
#-*- coding: utf-8 -*-

# Author: Edward Roualdes
#         [2012.11.20]
#         University of Kentucky

from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
from Bio.Seq import Seq
from Bio.Blast import NCBIXML
from collections import defaultdict


E_THRESH = 2e-5                 # E_VALUE threshold

# function definitions
def readBlastXml(blastReportXml):
    "extracts only hits with positive results from BLASTREPORT and stores necessary info in a defaultdict with key = blastRecord.query (aka read.description), value = [E_VALUE (aka hsps[0].expect), Locus (blastRecord.alignment[0].hit_def)]"

    bi = defaultdict(list)        # initialize data structure
    xOpen = open(blastReportXml, "rU") # open file for reading
    blastRecords = NCBIXML.parse(xOpen) 
    
    def m(bRec):
        """if a read has positive hits, record best hit as judged by the E_VALUE."""
        r = bRec.alignments      
        description = str(bRec.query)
        if r:                   # if positive hit
            hsp = r[0].hsps[0]  # store call to member hsps
            if hsp.expect < E_THRESH: # if E_VALUE sufficiently small
                bi[description].append(str(r[0].hit_def)) # Locus; element 0
                bi[description].append(hsp.frame[1]) # seq direction; element 1
                bi[description].append(hsp.query) # the query sequence; element 2
    map(m, blastRecords)
    xOpen.close()

    print "Read", len(bi.items()), "BLAST hits..."
    return bi
    

def sortLoci(sortedBarcodes, blastReportXml, c):
    "sort sortedBarcodes by locus if E_VALUE < E_THRESH, otherwise declare match isn't good enough and discard read"

# input:
# sortedBarcodes        = defaultdict() data structure output from sortBarcode

# output:
# sortedBarcodesLoci    = defaultdict() w/ key = "Locus", val = locus in each read's annotations

    bInfo = readBlastXml(blastReportXml) # parse BLASTREPORT
    sortedLoci = defaultdict(list)       # new data strct to exclude non-positive hit seqs

    for inds, reads in sortedBarcodes.iteritems(): # for each individual 
        for read in reads:                         # and each read
            data = bInfo[read.description]
            if data:
                tmp = SeqRecord(Seq(data[2])) # copy of hsp.query; use BLAST hit query
                tmp.id = read.id              # and copy over other pertinent info
                tmp.description = read.description
                tmp.name = read.name
                tmp.annotations["Locus"] = data[0] # annotate tmp with locus
                if data[1] == 1:                    # reversed hit from BLAST?
                    sortedLoci[inds].append(tmp)
                else:
                    sortedLoci[inds].append(
                        tmp.reverse_complement(
                            id = tmp.id,
                            name = tmp.name,
                            description = tmp.description,
                            annotations = tmp.annotations))            
            else:
                c -= 1
                    
    del sortedBarcodes, bInfo
    print "Classified", c, "reads by loci..."
    return sortedLoci
                        
                        
# test function: $ python sortLoci.py [INPUTFILE] [BARCODEFILE] [BLASTREPORT]

if __name__ == "__main__":
    # import sys
    from sortBarcode import sortBarcode
    # from pprint import pprint

    # i = sys.argv[1]             # INPUTFILE
    # bar = sys.argv[2]           # BARCODEFILE
    # bR = sys.argv[3]            # BLASTREPORT
    # c = 0                       # number of sorted reads

    # srt, c = sortBarcode(i, bar, True)
    # srt = sortLoci(srt, bR, c)
